import torch  # 导入PyTorch，用于构建和训练神经网络模型
import torch.nn as nn  # 导入nn模块，用于定义神经网络的各种层和损失函数

X = torch.tensor([-1.0, 0.0, 1.0, 2.0, 3.0, 4.0]).view(-1,1)  # 定义输入数据X，改变结构为了之后前向传播使用
y = torch.tensor([-3.0, -1.0, 1.0, 3.0, 5.0, 7.0]).view(-1,1)  # 定义真实标签y，改变结构为了之后前向传播使用

# 定义感知机模型，继承自nn.Module
class Perceptron(nn.Module):
    def __init__(self):
        super().__init__() # 调用父类的构造函数
        self.fc = nn.Linear(1,6) # 一个全连接层,输入和输出维度都是1
        self.fc = nn.Linear(6,1)

    # nn.Module类中的前向传播函数forward是一个抽象函数, 需要用户自己实现
    def forward(self, x):      
        y = self.fc.forward(x) # 调用nn.Linear类的forward函数, 可以简写成 y = self.fc(x)
        return y

model = Perceptron() # 实例化一个感知机模型
print("已经创建模型:", model) 
if model.training == True:
    print("模型当前(默认)处于训练模式")

loss_fnc = nn.MSELoss(reduction='mean')  # 定义损失函数为均方误差
optimizer = torch.optim.SGD(model.parameters(), lr=0.01) # 定义优化器为随机梯度下降, 学习率为0.01

# 循环500次，每次都进行一次训练
for epoch in range(500):
    # 使用模型对输入X进行预测，得到预测结果y_pred
    y_pred = model.forward(X)
    # 计算预测结果y_pred和真实值y之间的损失
    loss = loss_fnc(y_pred, y)
    
    # 清零优化器中的梯度信息，以便于进行下一轮的优化
    # 之所以要把这句话放在这里，是因为在PyTorch中，梯度是累加的，所以每次进行优化之前都要清零
    optimizer.zero_grad()
    
    # 对损失进行反向传播，计算梯度
    loss.backward()
    # 使用优化器对模型参数进行更新
    optimizer.step()

    # 每50次训练，打印一次当前的训练轮数和损失
    if (epoch+1) % 50 == 0:
        print(f'Epoch [{epoch+1}/500], Loss: {loss.item():.4f}')

model.eval() # 将模型调整为评估模式
if model.training == False:
    print("模型当前处于评估模式")

# no_grad()是一个上下文管理器，被该语句包裹的代码块将不会进行梯度计算
# 这样可以减少内存的占用，加速模型的预测
with torch.no_grad():
    x_test = torch.tensor([10.0])
    y_test = model.forward(x_test.unsqueeze(0))
    print(f'测试输入x: {x_test.item()}')  
    print(f'测试输出y: {y_test.item()}')